{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install gigachain faiss-cpu sentence-transformers sentencepiece rank_bm25 datasets --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "from langchain.retrievers import BM25Retriever, EnsembleRetriever\n",
    "from langchain.vectorstores.faiss import FAISS\n",
    "from langchain_core.documents import Document"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Данные"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Загрузим Question Answering датасет на русском языке - [SberQuAD](https://huggingface.co/datasets/sberquad).\n",
    "\n",
    "Каждая запись содержит несколько полей, из которых нам нужны будут только 2:\n",
    "- **question** - неструктурированный запрос или обычный вопрос;\n",
    "- **contex** - релевантный этому вопросу параграф из базы знаний (пассаж)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': 62310,\n",
       " 'title': 'SberChallenge',\n",
       " 'context': 'В протерозойских отложениях органические остатки встречаются намного чаще, чем в архейских. Они представлены известковыми выделениями сине-зелёных водорослей, ходами червей, остатками кишечнополостных. Кроме известковых водорослей, к числу древнейших растительных остатков относятся скопления графито-углистого вещества, образовавшегося в результате разложения Corycium enigmaticum. В кремнистых сланцах железорудной формации Канады найдены нитевидные водоросли, грибные нити и формы, близкие современным кокколитофоридам. В железистых кварцитах Северной Америки и Сибири обнаружены железистые продукты жизнедеятельности бактерий.',\n",
       " 'question': 'чем представлены органические остатки?',\n",
       " 'answers': {'text': ['известковыми выделениями сине-зелёных водорослей'],\n",
       "  'answer_start': [109]}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = datasets.load_dataset(\"sberquad\")\n",
    "ds[\"train\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_ds = ds[\"validation\"]\n",
    "documents = [\n",
    "    Document(page_content=context) for context in set(validation_ds[\"context\"])\n",
    "]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ретривал"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Наша задача - среди всех возможных пассажей найти наиболее подходящий вопросу, чтобы на основе него ГигаЧат смог дать максимально точный ответ.\n",
    "\n",
    "Для этой задачи существует множество обёрток различных retrieval подходов внутри пакета `langchain.retrievers`. Попробуем найти наиболее точный для наших данных"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Эмбеддинг модель"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Попробуем использовать нейросетевой ретривер (эмбеддер) с векторной базой данных.\n",
    "\n",
    "В качестве эмбеддингов возьмём специализированную мультиязычную модель [intfloat/multilingual-e5-base](https://huggingface.co/intfloat/multilingual-e5-base). Модель требует использовать различные префиксы для запросов и пассажей, поэтому создадим свой класс `HuggingFaceE5Embeddings` под эту логику, немного переопределив стандартный класс `langchain.embeddings.HuggingFaceEmbeddings`.\n",
    "\n",
    "Вы можете попробовать использовать любую энкодер модель для вашей задачи. Для русского языка существует открытый бенчмарк энкодеров предложений [avidale/encodechka](https://github.com/avidale/encodechka#лидерборд), который содержит отдельный лидерборд для ранжирования. Также можете посмотреть специализированный мультиязычный лидерборд [MTEB](https://huggingface.co/spaces/mteb/leaderboard)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Coroutine, List\n",
    "\n",
    "\n",
    "class HuggingFaceE5Embeddings(HuggingFaceEmbeddings):\n",
    "    def embed_query(self, text: str) -> List[float]:\n",
    "        text = f\"query: {text}\"\n",
    "        return super().embed_query(text)\n",
    "\n",
    "    def embed_documents(self, texts: List[str]) -> List[List[float]]:\n",
    "        texts = [f\"passage: {text}\" for text in texts]\n",
    "        return super().embed_documents(texts)\n",
    "\n",
    "    async def aembed_query(self, text: str) -> Coroutine[Any, Any, List[float]]:\n",
    "        text = f\"query: {text}\"\n",
    "        return await super().aembed_query(text)\n",
    "\n",
    "    async def aembed_documents(\n",
    "        self, texts: List[str]\n",
    "    ) -> Coroutine[Any, Any, List[List[float]]]:\n",
    "        texts = [f\"passage: {text}\" for text in texts]\n",
    "        return await super().aembed_documents(texts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding = HuggingFaceE5Embeddings(model_name=\"intfloat/multilingual-e5-base\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "А в качестве векторной базы данных используем [faiss](https://faiss.ai/), как самое популярное и простое решение для подобных задач. При создании экземпляра векторной БД для каждого документа считается его векторное представление, поиск по которому и происходит при каждом запросе - вся база данных ранжируется по степени релевантности входному запросу. В результате `embedding_retriever` вернёт `k = 5` самых подходящих пассажей.\n",
    "\n",
    "GigaChain, как и Langchain, предоставляет большое количество интерфейсов под все самые популярные векторные БД, полный список можно найти на [странице документации](https://python.langchain.com/docs/integrations/vectorstores)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "faiss_db = FAISS.from_documents(documents, embedding=embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_retriever = faiss_db.as_retriever(search_kwargs={\"k\": 5})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 5036/5036 [08:03<00:00, 10.42 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['id', 'title', 'context', 'question', 'answers', 'embedding_retrieved'],\n",
       "    num_rows: 5036\n",
       "})"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_ds = validation_ds.map(\n",
    "    lambda x: {\n",
    "        \"embedding_retrieved\": [\n",
    "            passage.page_content\n",
    "            for passage in embedding_retriever.get_relevant_documents(x[\"question\"])\n",
    "        ]\n",
    "    }\n",
    ")\n",
    "validation_ds"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В качестве целевой метрики используем точность нахождения целевого пассажа среди топ-5 отранжированных"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def acc_top(dataset: datasets.Dataset, right_column: str, answer_column: str) -> float:\n",
    "    temp_dataset = dataset.map(\n",
    "        lambda x: {\"is_right_retrieved\": x[right_column] in x[answer_column]}\n",
    "    )\n",
    "    return sum(temp_dataset[\"is_right_retrieved\"]) / len(temp_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 5036/5036 [00:01<00:00, 3481.41 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.9116362192216044"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "acc_top(validation_ds, \"context\", \"embedding_retrieved\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BM25"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Теперь попробуем оценить качество другого ретривала - [BM25](https://ru.wikipedia.org/wiki/Okapi_BM25)\n",
    "\n",
    "Попробуйте использовать токенизацию предложений на русском от библиотеки [natasha/razdel](https://github.com/natasha/razdel#usage), [aatimofeev/spacy_russian_tokenizer](https://github.com/aatimofeev/spacy_russian_tokenizer#usage) или [Koziev/rutokenizer](https://github.com/Koziev/rutokenizer#примеры) - они могут улучшить точность поиска на ваших данных"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import string\n",
    "\n",
    "\n",
    "def tokenize(s: str) -> list[str]:\n",
    "    \"\"\"Очень простая функция разбития предложения на слова\"\"\"\n",
    "    return s.lower().translate(str.maketrans(\"\", \"\", string.punctuation)).split(\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "bm25_retriever = BM25Retriever.from_documents(\n",
    "    documents=documents,\n",
    "    preprocess_func=tokenize,\n",
    "    k=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 5036/5036 [06:34<00:00, 12.75 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['id', 'title', 'context', 'question', 'answers', 'embedding_retrieved', 'bm25_retrieved'],\n",
       "    num_rows: 5036\n",
       "})"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_ds = validation_ds.map(\n",
    "    lambda x: {\n",
    "        \"bm25_retrieved\": [\n",
    "            passage.page_content\n",
    "            for passage in bm25_retriever.get_relevant_documents(x[\"question\"])\n",
    "        ]\n",
    "    }\n",
    ")\n",
    "validation_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 5036/5036 [00:01<00:00, 3617.90 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.9197776012708498"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "acc_top(validation_ds, \"context\", \"bm25_retrieved\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ансамбль"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "А теперь попробуем увеличить точность поиска, объединив работу двух предыдущих ретривалов алгоритмом [Reciprocal Rank Fusion](https://plg.uwaterloo.ca/~gvcormac/cormacksigir09-rrf.pdf). Для него есть удобная обёртка `langchain.retrievers.EnsembleRetriever`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_retriever = faiss_db.as_retriever(search_kwargs={\"k\": 2})\n",
    "bm25_retriever = BM25Retriever.from_documents(\n",
    "    documents=documents,\n",
    "    preprocess_func=tokenize,\n",
    "    k=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "ensemble_retriever = EnsembleRetriever(\n",
    "    retrievers=[embedding_retriever, bm25_retriever],\n",
    "    weights=[0.4, 0.6],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 5036/5036 [20:28<00:00,  4.10 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['id', 'title', 'context', 'question', 'answers', 'embedding_retrieved', 'bm25_retrieved', 'retrieved'],\n",
       "    num_rows: 5036\n",
       "})"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_ds = validation_ds.map(\n",
    "    lambda x: {\n",
    "        \"retrieved\": [\n",
    "            passage.page_content\n",
    "            for passage in ensemble_retriever.get_relevant_documents(x[\"question\"])\n",
    "        ]\n",
    "    }\n",
    ")\n",
    "validation_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|██████████| 5036/5036 [00:01<00:00, 3121.68 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.9690230341540905"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "acc_top(validation_ds, \"context\", \"retrieved\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## E2E решение"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "На улучшенном решении уже можно строить вопросно-ответную систему с достаточно хорошей точностью ответов. Ограничимся простой цепочкой `langchain.chains.RetrievalQA`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "from langchain.llms.gigachat import GigaChat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = GigaChat(credentials=...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "qa = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    chain_type=\"stuff\",\n",
    "    retriever=ensemble_retriever,\n",
    "    return_source_documents=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'query': 'Что такое вода?',\n",
       " 'result': 'Вода — это бинарное неорганическое соединение, состоящее из двух атомов водорода и одного атома кислорода, соединенных ковалентной связью. Она является прозрачной жидкостью без цвета, запаха и вкуса при нормальных условиях. Вода может существовать в различных состояниях, включая жидкое, твердое и газообразное. Она играет важную роль в жизни на Земле и является основой для существования всех живых организмов.',\n",
       " 'source_documents': [Document(page_content='На самом деле из-за низкого давления вода не может существовать в жидком состоянии на большей части (около 70 %) поверхности Марса. Вода в состоянии льда была обнаружена в марсианском грунте космическим аппаратом НАСА Феникс . В то же время собранные марсоходами Спирит и Opportunity геологические данные позволяют предположить, что в далёком прошлом вода покрывала значительную часть поверхности Марса. Наблюдения в течение последнего десятилетия позволили обнаружить в некоторых местах на поверхности Марса слабую гейзерную активность. По наблюдениям с космического аппарата Mars Global Surveyor , некоторые части южной полярной шапки Марса постепенно отступают.'),\n",
       "  Document(page_content='Оптические свойства воды оцениваются по её прозрачности, которая в свою очередь зависит от длины волны излучения, проходящего через воду. Вследствие поглощения оранжевых и красных компонентов света вода приобретает голубоватую окраску. Вода прозрачна только для видимого света и сильно поглощает инфракрасное излучение, поэтому на инфракрасных фотографиях водная поверхность всегда получается чёрной. Ультрафиолетовые лучи легко проходят через воду, поэтому растительные организмы способны развиваться в толще воды и на дне водоёмов, инфракрасные лучи проникают только в поверхностный слой. Вода отражает 5 % солнечных лучей, в то время как снег — около 85 %. Под лёд океана проникает только 2 % солнечного света.'),\n",
       "  Document(page_content='Концентрация инертных газов, аргона, криптона и ксенона, превышает их количество на Солнце (см. таблицу), а концентрация неона явно меньше. Присутствует незначительное количество простых углеводородов: этана, ацетилена и диацетилена, — которые формируются под воздействием солнечной ультрафиолетовой радиации и заряженных частиц, прибывающих из магнитосферы Юпитера. Диоксид углерода, моноксид углерода и вода в верхней части атмосферы, как полагают, своим присутствием обязаны столкновениям с атмосферой Юпитера комет, таких, например, как комета Шумейкеров-Леви 9. Вода не может прибывать из тропосферы, потому что тропопауза, действующая как холодная ловушка, эффективно препятствует поднятию воды до уровня стратосферы.'),\n",
       "  Document(page_content='Вода́ (оксид водорода) — бинарное неорганическое соединение с химической формулой Н2O. Молекула воды состоит из двух атомов водорода и одного — кислорода, которые соединены между собой ковалентной связью. При нормальных условиях представляет собой прозрачную жидкость, не имеющую цвета (при малой толщине слоя), запаха и вкуса. В твёрдом состоянии называется льдом (кристаллы льда могут образовывать снег или иней), а в газообразном — водяным паром. Вода также может существовать в виде жидких кристаллов (на гидрофильных поверхностях). Составляет приблизительно около 0,05 % массы Земли.'),\n",
       "  Document(page_content='друг другу; друг (о, в) друге; один (у, за, на, из, из-под, для) другого; друг (у, за, перед) дружкой; друг (у, за, на, из, из-под, для) друга; друг (с, за, над, под, перед) другом; друг (о, в) друге; один (у, за, на, из, для) другого; один (в, за, на) один; один к одному (другому); один (в, за, на) один; друг (с, за, под, перед) дружкой; друг (у, из, из-под) дружки; друг на дружке; раз за (на) раз[ом]; от раза к разу; раз к разу; от случая к случаю; каждый (у, за, на, из, для) каждого; каждый за (над, под, перед) каждым. каждый в каждом; тот (у, в, за, на, из, из-под, для) [э]того; от того к [э]тому; в конце концов; от начала к началу; от первого ко второму; от противного к противному;')]}"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qa.invoke({\"query\": \"Что такое вода?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
